package raft

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/rpc"
	"raft-go/raft/storage"
	"sync"
	"time"
)

type RPCProxy struct {
	mu                 sync.Mutex
	cm                 *ConsensusModule
	numCallsBeforeDrop int
}

type Server struct {
	mu sync.Mutex

	serverId int
	peerIds  []int

	cm       *ConsensusModule
	rpcProxy *RPCProxy

	rpcServer *rpc.Server
	listener  net.Listener

	storage storage.Storage

	commitChan  chan<- CommitEntry
	peerClients map[int]*rpc.Client

	ready <-chan any
	quit  chan any
	wg    sync.WaitGroup
}

func NewServer(id int, peers []int, storage storage.Storage, ready <-chan any, commitChan chan<- CommitEntry) *Server {
	server := new(Server)
	server.serverId = id
	server.peerIds = peers
	server.peerClients = make(map[int]*rpc.Client)
	server.ready = ready
	server.storage = storage
	server.commitChan = commitChan
	server.quit = make(chan any)
	return server
}

func (s *Server) Serve() {
	s.mu.Lock()
	s.rpcServer = rpc.NewServer()
	s.cm = NewConsensusModule(s.serverId, s.peerIds, s, s.storage, s.ready, s.commitChan)
	s.rpcProxy = &RPCProxy{cm: s.cm}
	var err error
	s.listener, err = net.Listen("tcp", ":0")
	if err != nil {
		log.Fatal("Listen TCP port failed...")
	}
	log.Printf("[%d] listening at %s\n", s.serverId, s.listener.Addr().String())
	s.rpcServer.RegisterName("ConsensusModule", s.rpcProxy)
	s.mu.Unlock()

	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		for {
			conn, err := s.listener.Accept()
			if err != nil {
				select {
				case <-s.quit:
					return
				default:
					log.Fatalf("build connection failed\n")
				}
			}
			s.wg.Add(1)
			go func() {
				s.rpcServer.ServeConn(conn)
				s.wg.Done()
			}()
		}
	}()
}

func (s *Server) DisconnectAll() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, id := range s.peerIds {
		if s.peerClients[id] != nil {
			s.peerClients[id].Close()
			s.peerClients[id] = nil
		}
	}
}

func (s *Server) Shutdown() {
	s.cm.Stop()
	close(s.quit)
	s.listener.Close()
	s.wg.Wait()
}

func (s *Server) GetListenAddr() net.Addr {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.listener.Addr()
}

func (s *Server) ConnectToPeer(peerid int, addr net.Addr) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 如果还未与其他节点建立连接，则新建客户端，并保存连接
	if s.peerClients[peerid] == nil {
		client, err := rpc.Dial(addr.Network(), addr.String())
		if err != nil {
			return err
		}
		s.peerClients[peerid] = client
	}
	return nil
}

func (s *Server) DisconnectToPeer(peerid int) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.peerClients[peerid] != nil {
		err := s.peerClients[peerid].Close()
		s.peerClients[peerid] = nil
		return err
	}
	return nil
}

func (s *Server) Submit(cmd any) int {
	return s.cm.Submit(cmd)
}

func (s *Server) Proxy() *RPCProxy {
	return s.rpcProxy
}

func (s *Server) Call(id int, method string, args any, reply any) error {
	s.mu.Lock()
	peerClient := s.peerClients[id]
	s.mu.Unlock()
	if peerClient == nil {
		return fmt.Errorf("peer client %d is disconnected", id)
	}
	return peerClient.Call(method, args, reply)
}

func (rpc *RPCProxy) RequestVote(args RequestVoteArgs, reply *RequestVoteReply) error {
	time.Sleep(time.Duration(1+rand.Intn(5)) * time.Millisecond)
	return rpc.cm.RequestVote(args, reply)
}

func (rpc *RPCProxy) AppendEntries(args AppendEnriesArgs, reply *AppendEntriesReply) error {
	time.Sleep(time.Duration(1+rand.Intn(5)) * time.Millisecond)
	return rpc.cm.AppendEntries(args, reply)
}

func (rpp *RPCProxy) DropCallsAfterN(n int) {
	rpp.mu.Lock()
	defer rpp.mu.Unlock()

	rpp.numCallsBeforeDrop = n
}

func (rpp *RPCProxy) DontDropCalls() {
	rpp.mu.Lock()
	defer rpp.mu.Unlock()

	rpp.numCallsBeforeDrop = -1
}
